import torch
from torch import nn
lambda_ = 1
class AutoRec(nn.Module):
    def __init__(self, m):
        super(AutoRec, self).__init__()
        self.m = m
        self._encoder = nn.Sequential(
            nn.Linear(self.m, 7),
            nn.Sigmoid()
        )
        self._decoder = nn.Sequential(
            nn.Linear(7, self.m),
            #nn.Sigmoid()
        )

    def forward(self, input):
        return self._decoder(self._encoder(input))



if __name__ == "__main__":
    m = 4
    n = 4
    data = torch.Tensor([[1, 2, 3, 4],
                         [5, 6, 7, 8],
                         [9, 10, 11, 12],
                         [13, 14, 15, 16]])

    c = torch.gather(data, dim=1, index=torch.LongTensor([[3 for k in range(m)]]).t())

    model = AutoRec(m)



    optim = torch.optim.SGD(model.parameters(), lr=0.01)
    loss_mse = torch.nn.MSELoss()






    for i in range(200):
        running_loss = 0.0
        output = model(c.t())
        result_mse = loss_mse(c.t(), output) * m
        cost = 0
        for group in optim.param_groups:
            for param in group['params']:
                if param.data.dim() == 2:
                    cost += torch.t(param.data).pow(2).sum()
        cost = cost * lambda_ * 0.5 + result_mse
        optim.zero_grad()
        cost.backward()
        optim.step()
        running_loss = running_loss + cost
        print(running_loss)